package com.example.langchainrag.config;


import com.example.langchainrag.service.Assistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
@RequiredArgsConstructor
public class AssistantConfig {
    final ChatLanguageModel chatLanguageModel;

//    @Bean
//    public EmbeddingStore<TextSegment> initEmbeddingStore() {
//        return new InMemoryEmbeddingStore<>();
//    }

    @Bean
    public Assistant assistant(SearchApiWebSearchEngine engine, EmbeddingStore<TextSegment> embeddingStore){
        return AiServices.builder(Assistant.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(15))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .tools(new WebSearchTool(engine))
                .chatLanguageModel(chatLanguageModel)
                .build();
    }
}
